import json, os
from .chess import flip_boardkey

_repr = {'Pawn':'P','Bishop':'B','Guard':'G','Cannon':'C',
		'Knight':'N','Rock':'R','King':'K'}
def board_key_repr(board_key):
	return ''.join([f'{c[0]}{_repr[t]}{i}{j}' for c,t,i,j in board_key])

class Cache:
	def __init__(self, cache_path=None):
		self._cache = {}
		if cache_path is not None:
			if os.path.isfile(cache_path):
				print(f'读取缓存 {cache_path}')
				self.load(cache_path)
			def save():
				print(f'写入缓存 {cache_path}')
				self.save(cache_path)
			import atexit
			atexit.register(save)

	def _score_from_cache(self, node, board_key):
		key = board_key_repr(board_key)
		if key in self._cache:
			height, score = self._cache[key]
			if height >= node.height:
				node.score = score

	def load(self, path):
		with open(path) as fp:
			self._cache = json.load(fp)

	def save(self, path):
		with open(path, 'w') as fp:
			json.dump(self._cache, fp, indent='\t')

	def add_to_cache(self, node):
		key = board_key_repr(node.board_key)
		if key not in self._cache:
			key = board_key_repr(flip_boardkey(node.board_key))
		if key in self._cache:
			if self._cache[key][0] > node.height:
				return
		self._cache[key] = (node.height, node.score)

	def score_from_cache(self, node, parent):
		self._score_from_cache(node, node.board_key)
		self._score_from_cache(node, flip_boardkey(node.board_key))
		if node.score is not None:
			self.add_to_cache(node)


class AlphaBetaCache(Cache):
	def _score_from_cache(self, node, board_key):
		key = board_key_repr(board_key)
		if key in self._cache:
			for height, score, over_beta in self._cache[key]:
				if height<node.height: break
				if over_beta is None:
					node.score = score
					break
				assert isinstance(over_beta, bool), over_beta
				if over_beta is True: # score>=prev_beta, score是下界,score* >=score
					if score >= node.beta:
						node.score = score
						break
				else: # score<=prev_alpha, score是上界,score* <=score
					if score <= node.alpha:
						node.score = score
						break

	def add_to_cache(self, node):
		over_beta = None
		if not (node.alpha<node.score<node.beta):
			over_beta = (node.score>=node.beta)
		self._add_to_cache(node.board_key, node.height, node.score, over_beta)

	def _add_to_cache(self, board_key_, height_, score_, over_beta_):
		key = board_key_repr(board_key_)
		if key not in self._cache:
			key = board_key_repr(flip_boardkey(board_key_))
		if key not in self._cache:
			self._cache[key] = []
		info = (height_, score_, over_beta_)
		cache = self._cache[key]
		def substitute(height1, score1, over_beta1, height2, score2, over_beta2): 
			assert height1 >= height2
			if over_beta1 is None: return True
			if score1 < score2:
				return (over_beta1 is False) and (over_beta2 is False)
			else: # score1 >= score2 
				return (over_beta1 is True) and (over_beta2 is True)
		updated = False
		for i, (height, score, over_beta) in enumerate(cache):
			if height_ >= height:
				if substitute(height_, score_, over_beta_, height, score, over_beta):
					cache[i] = None
					updated = True
		if updated:
			cache = [c for c in cache if c is not None]
		for i, (height, score, over_beta) in enumerate(cache):
			if height_ <= height:
				if substitute(height, score, over_beta, height_, score_, over_beta_):
					info = None
					break
		if updated or info is not None:
			if info is not None:
				cache.append(info)
				cache.sort(key=lambda c: c[0], reverse=True) # 按height从大到小排
			self._cache[key] = cache

	def score_from_cache(self, node, parent):
		node.alpha, node.beta = -parent.beta, -parent.alpha
		super().score_from_cache(node, parent)

	def merge(self, other):
		_repr1 = {'R':'Red','B':'Black'}
		_repr2 = {v:k for k,v in _repr.items()}
		from tqdm import tqdm
		for board_key, info in tqdm(other._cache.items(), total=len(other._cache)):
			board_key = [board_key[i*4:i*4+4] for i in range(len(board_key)//4)]
			board_key = tuple((_repr1[k[0]],_repr2[k[1]],int(k[2]),int(k[3])) for k in board_key)
			for height_, score_, over_beta_ in info:
				self._add_to_cache(board_key, height_, score_, over_beta_)

if __name__ == '__main__':
	import gc
	c = AlphaBetaCache ()
	for i in range(2):
		c1 = AlphaBetaCache()
		filename = f'{i+1}/cache.json'
		print(filename)
		c1.load(filename)
		print('merging')
		c.merge(c1)
		del c1
		print('gc')
		gc.collect()
	c.save('merged_cache.json')

